package com.sgj.codecheck;

import com.intellij.openapi.actionSystem.AnAction;
import com.intellij.openapi.actionSystem.AnActionEvent;
import com.intellij.openapi.application.ApplicationManager;
import com.intellij.openapi.application.ReadAction;
import com.intellij.openapi.compiler.CompileScope;
import com.intellij.openapi.compiler.CompilerManager;
import com.intellij.openapi.diagnostic.Logger;
import com.intellij.openapi.progress.ProgressIndicator;
import com.intellij.openapi.progress.ProgressManager;
import com.intellij.openapi.progress.Task;
import com.intellij.openapi.project.Project;
import com.intellij.openapi.ui.DialogBuilder;
import com.intellij.openapi.ui.DialogWrapper;
import com.intellij.openapi.ui.Messages;
import com.intellij.openapi.vfs.VirtualFile;
import com.intellij.ui.components.JBList;
import org.jetbrains.annotations.NotNull;

import javax.swing.*;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
import java.util.stream.Collectors;

/**
 * @author GuojiShen
 * @date 2025/7/24 15:55
 * @since 2025-07-24
 */
public class CodeQlCheck extends AnAction {

    private static final Logger LOG = Logger.getInstance(CodeQlCheck.class);
    private static final String CODEQL_PATH = "D:\\Tool\\codeql\\codeql.exe";
    private static final String CODEQL_QUERY_DIR = "D:\\Tool\\codeql-full\\java\\ql\\src\\Security\\CWE";

    @Override
    public void actionPerformed(@NotNull AnActionEvent e) {
        Project project = e.getProject();
        if (project == null) {
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showErrorDialog("No project found.", "Error"));
            return;
        }

        // 弹出对话框让用户输入路径
        String path = Messages.showInputDialog(project, "请输入需要扫描的根路径:",
                "扫描路径", Messages.getQuestionIcon(), "", null);
        if (path == null || path.trim().isEmpty()) {
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showWarningDialog(project, "路径不能为空", "Warning"));
            return;
        }

        // 弹出对话框选择查询文件
        List<String> selectedQueries = showQuerySelectionDialog(project);
        if (selectedQueries.isEmpty()) {
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showWarningDialog(project, "未选择任何查询文件", "Warning"));
            return;
        }

        // 异步执行扫描任务
        ProgressManager.getInstance().run(new Task.Backgroundable(project, "Scanning Files with CodeQL", true) {
            @Override
            public void run(@NotNull ProgressIndicator indicator) {
                indicator.setIndeterminate(false);
                indicator.setFraction(0.1);
                try {
                    indicator.setText("Scanning for Java files...");
                    List<VirtualFile> files = scanPath(project, path);
                    if (files.isEmpty()) {
                        ApplicationManager.getApplication().invokeLater(() ->
                                Messages.showInfoMessage(project, "No Java files found in the specified path.", "Scan Result"));
                        return;
                    }

                    indicator.setFraction(0.3);
                    ApplicationManager.getApplication().invokeLater(() ->
                            Messages.showInfoMessage(project, "Found " + files.size() + " Java files.", "Scan Result"));
                    files.forEach(file -> LOG.info("Found file: " + file.getPath()));

                    scanWithCodeQL(project, files, selectedQueries, indicator);
                } catch (Exception ex) {
                    LOG.error("Scan failed: ", ex);
                    ApplicationManager.getApplication().invokeLater(() ->
                            Messages.showErrorDialog(project, "Scan failed: " + ex.getMessage(), "Error"));
                }
            }
        });
    }

    private List<String> showQuerySelectionDialog(Project project) {
        File queryDir = new File(CODEQL_QUERY_DIR);
        if (!queryDir.exists() || !queryDir.isDirectory()) {
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showErrorDialog(project, "CodeQL query directory not found: " + CODEQL_QUERY_DIR, "Error"));
            return new ArrayList<>();
        }

        // 递归收集所有 .ql 文件
        List<String> queryPaths = new ArrayList<>();
        collectQueryFilesRecursively(queryDir, queryPaths);
        if (queryPaths.isEmpty()) {
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showErrorDialog(project, "No .ql files found in: " + CODEQL_QUERY_DIR, "Error"));
            return new ArrayList<>();
        }

        // 创建多选列表
        JBList<String> queryList = new JBList<>(queryPaths);
        queryList.setSelectionMode(ListSelectionModel.MULTIPLE_INTERVAL_SELECTION);

        // 使用 DialogBuilder 显示对话框
        DialogBuilder builder = new DialogBuilder(project);
        builder.setTitle("选择 CodeQL 查询文件");
        builder.setCenterPanel(new JScrollPane(queryList));
        builder.addOkAction();
        builder.addCancelAction();
        builder.setOkOperation(() -> builder.getDialogWrapper().close(DialogWrapper.OK_EXIT_CODE));

        boolean ok = builder.show() == DialogWrapper.OK_EXIT_CODE;
        if (!ok) {
            return new ArrayList<>();
        }
        return queryList.getSelectedValuesList();
    }

    private void collectQueryFilesRecursively(File dir, List<String> queryPaths) {
        File[] files = dir.listFiles();
        if (files != null) {
            for (File file : files) {
                if (file.isDirectory()) {
                    collectQueryFilesRecursively(file, queryPaths);
                } else if (file.getName().endsWith(".ql")) {
                    queryPaths.add(file.getAbsolutePath());
                }
            }
        }
    }

    private List<VirtualFile> scanPath(Project project, String path) {
        return ReadAction.compute(() -> {
            VirtualFile baseDir = project.getBaseDir();
            if (baseDir == null) {
                LOG.warn("Project base directory is null.");
                return new ArrayList<>();
            }

            VirtualFile targetDir = baseDir.findFileByRelativePath(path);
            if (targetDir == null || !targetDir.isDirectory()) {
                LOG.warn("Invalid or non-directory path: " + path);
                return new ArrayList<>();
            }

            List<VirtualFile> files = new ArrayList<>();
            collectFilesRecursively(targetDir, files);
            return files;
        });
    }

    private void collectFilesRecursively(VirtualFile dir, List<VirtualFile> files) {
        VirtualFile[] children = dir.getChildren();
        if (children != null) {
            for (VirtualFile child : children) {
                if (child.isDirectory()) {
                    collectFilesRecursively(child, files);
                } else if ("java".equalsIgnoreCase(child.getExtension())) {
                    files.add(child);
                }
            }
        }
    }

    private void compileProject(Project project, List<VirtualFile> files, Runnable onSuccess, ProgressIndicator indicator) {
        ApplicationManager.getApplication().invokeAndWait(() -> {
            CompilerManager compilerManager = CompilerManager.getInstance(project);
            CompileScope scope = compilerManager.createFilesCompileScope(files.toArray(new VirtualFile[0]));
            compilerManager.compile(scope, (aborted, errors, warnings, context) -> {
                if (!aborted && errors == 0) {
                    indicator.setText("Compilation successful. Running CodeQL analysis...");
                    indicator.setFraction(0.5);
                    onSuccess.run();
                } else {
                    ApplicationManager.getApplication().invokeLater(() ->
                            Messages.showErrorDialog(project, "Compilation failed with " + errors + " errors.", "Error"));
                }
            });
        });
    }

    private void scanWithCodeQL(Project project, List<VirtualFile> files, List<String> queryPaths, ProgressIndicator indicator) {
        String basePath = ReadAction.compute(project::getBasePath);
        if (basePath == null) {
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showErrorDialog(project, "Project base path is null.", "Error"));
            return;
        }

        File tempDir = new File(basePath, "codeql-temp");
        if (!tempDir.exists() && !tempDir.mkdirs()) {
            LOG.error("Failed to create temp directory: " + tempDir.getAbsolutePath());
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showErrorDialog(project, "Failed to create temp directory.", "Error"));
            return;
        }

        File codeqlExe = new File(CODEQL_PATH);
        if (!codeqlExe.exists()) {
            LOG.error("CodeQL executable not found: " + codeqlExe.getAbsolutePath());
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showErrorDialog(project, "CodeQL executable not found: " + codeqlExe.getAbsolutePath(), "Error"));
            return;
        }
        if (!codeqlExe.canExecute()) {
            LOG.error("CodeQL executable is not executable: " + codeqlExe.getAbsolutePath());
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showErrorDialog(project, "CodeQL executable is not executable: " + codeqlExe.getAbsolutePath(), "Error"));
            return;
        }

        for (String queryPath : queryPaths) {
            File queryFile = new File(queryPath);
            if (!queryFile.exists()) {
                LOG.error("Query file not found: " + queryPath);
                ApplicationManager.getApplication().invokeLater(() ->
                        Messages.showErrorDialog(project, "Query file not found: " + queryPath, "Error"));
                return;
            }
        }

        String dbPath = new File(tempDir, "project-db").getAbsolutePath();
        String sourceRoot = basePath;
        try {
            // 删除旧数据库
            File dbDir = new File(dbPath);
            if (dbDir.exists()) {
                deleteDirectory(dbDir);
                LOG.info("Deleted existing database: " + dbPath);
            }

            // 编译项目
            compileProject(project, files, () -> {
                try {
                    indicator.setText("Creating CodeQL database...");
                    indicator.setFraction(0.6);
                    // 设置 JAVA_HOME 为 Java 1.8
                    ProcessBuilder dbInit = new ProcessBuilder(
                            CODEQL_PATH, "database", "create", dbPath,
                            "--language=java", "--source-root=" + sourceRoot,
                            "--command=mvn clean package -B -DskipTests -Djava.version=1.8"
                    );
                    dbInit.directory(new File(CODEQL_PATH).getParentFile());
                    dbInit.redirectErrorStream(true);
                    Process dbInitProcess = dbInit.start();
                    String dbInitOutput = readProcessOutput(dbInitProcess);
                    int dbInitExitCode = dbInitProcess.waitFor();
                    if (dbInitExitCode != 0) {
                        LOG.error("Database creation failed: " + dbInitOutput);
                        throw new IOException("Database creation failed with exit code " + dbInitExitCode + ": " + dbInitOutput);
                    }
                    LOG.info("Database created successfully: " + dbPath);

                    // 运行查询并生成 CSV
                    indicator.setText("Running CodeQL queries...");
                    indicator.setFraction(0.8);
                    String outputCsvPath = new File(tempDir, "result.csv").getAbsolutePath();
                    List<String> command = new ArrayList<>();
                    command.add(CODEQL_PATH);
                    command.add("database");
                    command.add("analyze");
                    command.add(dbPath);
                    command.add("--format=csv");
                    command.add("--output=" + outputCsvPath);
                    command.addAll(queryPaths);
                    ProcessBuilder queryRun = new ProcessBuilder(command);
                    queryRun.directory(new File(CODEQL_PATH).getParentFile());
                    queryRun.redirectErrorStream(true);
                    Process queryProcess = queryRun.start();
                    String queryOutput = readProcessOutput(queryProcess);
                    int queryExitCode = queryProcess.waitFor();
                    if (queryExitCode != 0) {
                        LOG.error("Query execution failed: " + queryOutput);
                        throw new IOException("Query execution failed with exit code " + queryExitCode + ": " + queryOutput);
                    }

                    // 读取并处理 CSV 结果
                    indicator.setText("Processing CodeQL results...");
                    indicator.setFraction(0.9);
                    String csvResult = readFile(outputCsvPath);
                    processCsvResult(project, csvResult, files);
                    indicator.setFraction(1.0);
                } catch (IOException | InterruptedException ex) {
                    LOG.error("CodeQL scan failed: ", ex);
                    ApplicationManager.getApplication().invokeLater(() ->
                            Messages.showErrorDialog(project, "CodeQL scan failed: " + ex.getMessage(), "Error"));
                }
            }, indicator);
        } catch (Exception ex) {
            LOG.error("Scan setup failed: ", ex);
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showErrorDialog(project, "Scan setup failed: " + ex.getMessage(), "Error"));
        }
    }

    private String readFile(String path) throws IOException {
        try (BufferedReader reader = new BufferedReader(new FileReader(path, StandardCharsets.UTF_8))) {
            return reader.lines().collect(Collectors.joining("\n"));
        }
    }

    private String readProcessOutput(Process process) throws IOException {
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream(), StandardCharsets.UTF_8));
             BufferedReader errorReader = new BufferedReader(new InputStreamReader(process.getErrorStream(), StandardCharsets.UTF_8))) {
            String output = reader.lines().collect(Collectors.joining("\n"));
            String error = errorReader.lines().collect(Collectors.joining("\n"));
            return output + "\n" + error;
        }
    }

    private void processCsvResult(Project project, String csvOutput, List<VirtualFile> files) {
        if (csvOutput.isEmpty()) {
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showInfoMessage(project, "No issues found.", "CodeQL Scan Result"));
            return;
        }

        try {
            StringBuilder message = new StringBuilder("Found issues:\n");
            String[] lines = csvOutput.split("\n");
            boolean firstLine = true;
            for (String line : lines) {
                if (firstLine) {
                    firstLine = false; // 跳过 CSV 头部
                    continue;
                }
                if (line.trim().isEmpty()) {
                    continue;
                }
                // CSV 格式示例: "name,description,severity,message,path,startLine,startColumn,endLine,endColumn"
                String[] columns = line.split(",", -1); // -1 保留空字段
                if (columns.length >= 6) { // 确保至少有 name, description, severity, message, path, startLine
                    String name = columns[0].replace("\"", "");
                    String messageText = columns[3].replace("\"", "");
                    String path = columns[4].replace("\"", "");
                    String startLine = columns[5].replace("\"", "");
                    message.append(String.format("- %s in %s at line %s: %s\n", name, path, startLine, messageText));
                }
            }

            if (message.length() == "Found issues:\n".length()) {
                ApplicationManager.getApplication().invokeLater(() ->
                        Messages.showInfoMessage(project, "No issues found.", "CodeQL Scan Result"));
            } else {
                ApplicationManager.getApplication().invokeLater(() ->
                        Messages.showWarningDialog(project, message.toString(), "CodeQL Scan Result"));
            }
        } catch (Exception ex) {
            LOG.error("Failed to parse CSV output: ", ex);
            ApplicationManager.getApplication().invokeLater(() ->
                    Messages.showErrorDialog(project, "Failed to parse CodeQL results: " + ex.getMessage(), "Error"));
        }
    }

    private void deleteDirectory(File directory) {
        if (directory == null || !directory.exists()) {
            return;
        }
        if (directory.isDirectory()) {
            File[] files = directory.listFiles();
            if (files != null) {
                for (File file : files) {
                    if (file.isDirectory()) {
                        deleteDirectory(file);
                    } else {
                        if (!file.delete()) {
                            LOG.warn("Failed to delete file: " + file.getAbsolutePath());
                        }
                    }
                }
            }
        }
        if (!directory.delete()) {
            LOG.warn("Failed to delete directory: " + directory.getAbsolutePath());
        }
    }
}